import btk
import unittest
import _TDDConfigure
import numpy

class ForcePlatformsExtractorTest(unittest.TestCase):
    def test_FileSample01Eb015pi(self):
        reader = btk.btkAcquisitionFileReader()
        reader.SetFilename(_TDDConfigure.C3DFilePathIN + "sample01/Eb015pi.c3d")
        acq = reader.GetOutput()
        pfe = btk.btkForcePlatformsExtractor()
        pfe.SetInput(reader.GetOutput())
        pfc = pfe.GetOutput()
        pfe.Update()
        self.assertEqual(pfc.GetItemNumber(), 2)
        pf1 = pfc.GetItem(0)
        self.assertEqual(pf1.GetType(), 2)
        self.assertEqual(pf1.GetChannelNumber(), 6)
        self.assertEqual(pf1.GetChannel(0).GetLabel(), "FX1")
        self.assertEqual(pf1.GetChannel(1).GetLabel(), "FY1")
        self.assertEqual(pf1.GetChannel(2).GetLabel(), "FZ1")
        self.assertEqual(pf1.GetChannel(3).GetLabel(), "MX1")
        self.assertEqual(pf1.GetChannel(4).GetLabel(), "MY1")
        self.assertEqual(pf1.GetChannel(5).GetLabel(), "MZ1")
        o1 = numpy.array([[-4.4], [1.9], [-21.6]])
        numpy.testing.assert_array_almost_equal(pf1.GetOrigin(), o1)
        c11 = numpy.array([[520.0451], [1242.1694], [0.6219]])
        c21 = numpy.array([[57.0463], [1243.1996], [0.6211]])
        c31 = numpy.array([[58.1765], [1751.1963], [2.0812]])
        c41 = numpy.array([[521.1754], [1750.1661], [2.0820]])
        numpy.testing.assert_array_almost_equal(pf1.GetCorner(0), c11, 4)
        numpy.testing.assert_array_almost_equal(pf1.GetCorner(1), c21, 4)
        numpy.testing.assert_array_almost_equal(pf1.GetCorner(2), c31, 4)
        numpy.testing.assert_array_almost_equal(pf1.GetCorner(3), c41, 4)
        cs1 = numpy.concatenate((c11,c21,c31,c41), axis=1)
        numpy.testing.assert_array_almost_equal(pf1.GetCorners(), cs1, 4)
        pf2 = pfc.GetItem(1)
        self.assertEqual(pf2.GetType(), 2)
        self.assertEqual(pf2.GetChannelNumber(), 6)
        self.assertEqual(pf2.GetChannel(0).GetLabel(), "FX2")
        self.assertEqual(pf2.GetChannel(1).GetLabel(), "FY2")
        self.assertEqual(pf2.GetChannel(2).GetLabel(), "FZ2")
        self.assertEqual(pf2.GetChannel(3).GetLabel(), "MX2")
        self.assertEqual(pf2.GetChannel(4).GetLabel(), "MY2")
        self.assertEqual(pf2.GetChannel(5).GetLabel(), "MZ2")
        o2 = numpy.array([[-4.06] , [3.81], [-20.06]])
        numpy.testing.assert_array_almost_equal(pf2.GetOrigin(), o2)
        c12 = numpy.array([[53.6554]  , [1139.9977], [1.9204]])
        c22 = numpy.array([[516.6432] , [1143.3159], [1.2880]])
        c32 = numpy.array([[520.2825] , [635.3301], [0.1814]])
        c42 = numpy.array([[57.2948] , [632.0118], [0.8138]])
        numpy.testing.assert_array_almost_equal(pf2.GetCorner(0), c12, 4)
        numpy.testing.assert_array_almost_equal(pf2.GetCorner(1), c22, 4)
        numpy.testing.assert_array_almost_equal(pf2.GetCorner(2), c32, 4)
        numpy.testing.assert_array_almost_equal(pf2.GetCorner(3), c42, 4)
        cs2 = numpy.concatenate((c12,c22,c32,c42), axis=1)
        numpy.testing.assert_array_almost_equal(pf2.GetCorners(), cs2, 4)

    def test_FileSample10Type2(self):
        reader = btk.btkAcquisitionFileReader()
        reader.SetFilename(_TDDConfigure.C3DFilePathIN + "sample10/TYPE-2.C3D")
        pfe = btk.btkForcePlatformsExtractor()
        pfe.SetInput(reader.GetOutput())
        pfc = pfe.GetOutput()
        pfe.Update()
        self.assertEqual(pfc.GetItemNumber(), 1)
        pf1 = pfc.GetItem(0)
        self.assertEqual(pf1.GetType(), 2)
        self.assertEqual(pf1.GetChannelNumber(), 6)
        self.assertEqual(pf1.GetChannel(0).GetLabel(), "FX1")
        self.assertEqual(pf1.GetChannel(1).GetLabel(), "FY1")
        self.assertEqual(pf1.GetChannel(2).GetLabel(), "FZ1")
        self.assertEqual(pf1.GetChannel(3).GetLabel(), "MX1")
        self.assertEqual(pf1.GetChannel(4).GetLabel(), "MY1")
        self.assertEqual(pf1.GetChannel(5).GetLabel(), "MZ1")
        o1 = numpy.array([[-1.6] , [0.7], [37.5]])
        numpy.testing.assert_array_almost_equal(pf1.GetOrigin(), o1, 4)
        self.assertAlmostEqual(pf1.GetChannel(0).GetValues()[1], 0.08843, 5)
        self.assertAlmostEqual(pf1.GetChannel(0).GetValues()[3], -0.08843, 5)
        self.assertAlmostEqual(pf1.GetChannel(0).GetValues()[1020], -69.59441, 5)
        self.assertAlmostEqual(pf1.GetChannel(0).GetValues()[1021], -72.24731, 5)
        self.assertAlmostEqual(pf1.GetChannel(1).GetValues()[0], 0.17762, 5)
        self.assertAlmostEqual(pf1.GetChannel(3).GetValues()[2], 36.83129, 4)
        self.assertAlmostEqual(pf1.GetChannel(3).GetValues()[560], 73.66259, 4)
        self.assertAlmostEqual(pf1.GetChannel(3).GetValues()[561], 0.0, 5)

    def test_FileSample10Type4(self):
        reader = btk.btkAcquisitionFileReader()
        reader.SetFilename(_TDDConfigure.C3DFilePathIN + "sample10/TYPE-4.C3D")
        pfe = btk.btkForcePlatformsExtractor()
        pfe.SetInput(reader.GetOutput())
        pfc = pfe.GetOutput()
        pfe.Update()
        self.assertEqual(pfc.GetItemNumber(), 1)
        pf1 = pfc.GetItem(0)
        self.assertEqual(pf1.GetType(), 4)
        self.assertEqual(pf1.GetChannelNumber(), 6)
        self.assertEqual(pf1.GetChannel(0).GetFrameNumber(), 3980)
        self.assertEqual(pf1.GetChannel(0).GetLabel(), "FX1")
        self.assertEqual(pf1.GetChannel(1).GetLabel(), "FY1")
        self.assertEqual(pf1.GetChannel(2).GetLabel(), "FZ1")
        self.assertEqual(pf1.GetChannel(3).GetLabel(), "MX1")
        self.assertEqual(pf1.GetChannel(4).GetLabel(), "MY1")
        self.assertEqual(pf1.GetChannel(5).GetLabel(), "MZ1")
        cal = numpy.array([[ 1.5270, 0.0140, 0.0460,-4.2840,  -5.4940,  2.8560 ],
                           [-0.0050, 1.5390,-0.0190, 17.1600,  5.6680, -9.6390 ],
                           [ 0.0000,-0.0090, 5.9880, 13.2330, 13.2500,  4.6090 ],
                           [ 0.0000, 0.0020, 0.0000,741.8710, -0.5390, -0.5970 ],
                           [ 0.0020,-0.0010, 0.000,  -1.6820,739.6310,  2.0420 ],
                           [-0.0020,-0.0050,-0.0020, -3.2500, -0.5940,391.8790 ]])
        numpy.testing.assert_array_almost_equal(pf1.GetCalMatrix(), cal.transpose(), 4)
        o1 = numpy.array([[-1.6] , [0.7], [37.5]])
        numpy.testing.assert_array_almost_equal(pf1.GetOrigin(), o1)

    def test_FileSample19Sample19(self):
        reader = btk.btkAcquisitionFileReader()
        reader.SetFilename(_TDDConfigure.C3DFilePathIN + "sample19/sample19.c3d")
        pfe = btk.btkForcePlatformsExtractor()
        pfe.SetInput(reader.GetOutput())
        pfc = pfe.GetOutput()
        pfe.Update()
        self.assertEqual(pfc.GetItemNumber(), 4)
        pf1 = pfc.GetItem(0)
        self.assertEqual(pf1.GetChannelNumber(), 6)
        self.assertEqual(pf1.GetChannel(0).GetFrameNumber(), 34672 * 18)
        self.assertEqual(pf1.GetChannel(0).GetLabel(), "FP1C1")
        self.assertEqual(pf1.GetChannel(1).GetLabel(), "FP1C2")
        self.assertEqual(pf1.GetChannel(2).GetLabel(), "FP1C3")
        self.assertEqual(pf1.GetChannel(3).GetLabel(), "FP1C4")
        self.assertEqual(pf1.GetChannel(4).GetLabel(), "FP1C5")
        self.assertEqual(pf1.GetChannel(5).GetLabel(), "FP1C6")
        self.assertAlmostEqual(pf1.GetChannel(3).GetValues()[561], 0.0, 5)